from enum import unique
import os

from myapp.const.base import TEnum, defaultInt, defaultStr


# Defines the default deployment timeout when resources are scarce.
DEPLOYTIMEOUT_WHEN_LACK_OF_RESOURCE = 60
# Defines the maximum deployment timeout when resources are scarce, configurable via environment variable.
MAX_DEPLOYTIMEOUT_WHEN_LACK_OF_RESOURCE = int(
    os.getenv('MAX_DEPLOYTIMEOUT_WHEN_LACK_OF_RESOURCE', 5 * DEPLOYTIMEOUT_WHEN_LACK_OF_RESOURCE)
)
# Defines the Redis key used to store the deployment timeout setting when resources are scarce.
REDIS_KEY_DEPLOYTIMEOUT_WHEN_LACK_OF_RESOURCE = (
    'kubeflow-dashboard/deploy_timeout_when_lack_of_resource'
)


# 模型状态枚举
# Defines an enumeration for various model statuses.
class EnumModelStatus(TEnum):
    # Represents an undefined or unknown model status.
    undefined = defaultInt
    # Indicates that the model operation was successful.
    success = 1
    # Indicates that the model is currently being imported.
    importing = 2
    # Indicates that the model operation failed.
    fail = 3
    # Another representation for a failed model operation.
    failed = 3
    # Indicates that the model is currently being published.
    publishing = 2
    # Indicates that the model has been successfully published.
    published = 1

    @staticmethod
    # Retrieves the name of the model status given its value.
    # Args:
    #   val: The integer value of the model status.
    #   default: The default name to return if the value is not found. Defaults to 'undefined'.
    # Returns:
    #   The name of the model status as a string.
    def get_name(val, default=None):
        if default is None:
            default = EnumModelStatus.undefined.name
        try:
            key = EnumModelStatus(val).name
        except Exception:
            key = default
        return key


# 推理框架枚举
# Defines an enumeration for different serving frameworks.
@unique
class EnumServingFramework(TEnum):
    # Represents an unknown or other serving framework.
    other = defaultStr
    # Represents the MindSpore Serving framework.
    mindspore_serving = 'mindspore_serving'
    # Represents the TorchServe framework.
    torchserve = 'torchserve'
    # Represents the Taichu Serve framework.
    taichu_serve = 'taichu-serve'


# 算法框架枚举
# Defines an enumeration for different algorithm frameworks.
@unique
class EnumFramework(TEnum):
    # Represents an unknown or other algorithm framework.
    other = defaultStr
    # Represents the PyTorch framework.
    pytorch = 'pytorch'
    # Represents the MindSpore framework.
    mindspore = 'mindspore'


# 模型来源
# Defines an enumeration for different sources of models.
@unique
class EnumModelSource(TEnum):
    # Represents an unknown or other model source.
    other = defaultInt
    # Indicates that the model was created via auto-learning.
    auto_learning = 1  # 自动学习
    # Indicates that the model was created via custom training.
    training = 2  # 自定义训练
    # Indicates that the model was uploaded.
    upload = 3
    # Indicates that the model was imported from an image.
    image = 4  # 镜像导入
    # Indicates that the model was created via a visualization pipeline.
    pipeline = 5  # 可视化建模
    # Indicates that the model was created via large model finetuning.
    finetune = 6  # 大模型微调
    # Indicates that the model was created via incremental pre-training.
    post_pretraining = 7  # 增量预训练
    # Indicates that the model is an internal built-in model.
    internal = 8  # 内置模型
    # Indicates that the model was created via model compression.
    compress = 9  # 模型压缩
    # Indicates that the model was created via Reinforcement Learning from Human Feedback (RLHF).
    rlhf = 10 # 偏好对齐


# 模型类型
# Defines an enumeration for various types of models.
@unique
class EnumModelType(TEnum):
    # Represents an unknown or other model type.
    other = defaultInt
    # Model type for image captioning.
    image_caption = 1
    # Model type for visual question answering.
    vqa = 2
    # Model type for retrieval tasks.
    retrieval = 3
    # Model type for entity and relation extraction.
    entity_extraction = 4  # 实体关系抽取
    # Model type for poetry generation.
    poetry = 5  # 古诗生成
    # Model type for text-to-image generation.
    txt2image = 6  # 以文生图
    # Model type for multimodal emotion recognition.
    multimodal_emotion_recognition = 7  # 多模态情感识别
    # Model type for cross-modal retrieval (text-to-image, CLIP version).
    ret_txt_clip = 8  # 跨模态检索-文搜图 clip版本
    # Model type for virtual human lip animation.
    virtual_human_lip_animation = 9  # 虚拟人唇形动作
    # Model type for cartoon stylization of photos.
    cartoon_styled = 10  # 照片卡通风格化
    # Model type for face generation.
    face_generation = 11  # 人脸生成
    # Model type for dialogue bots.
    dialogue = 12  # 对话机器人
    # Model type for intent detection.
    intent_detect = 13  # 意图识别
    # Model type for text matching.
    text_match = 14  # 文本匹配
    # Model type for legal event detection.
    law_event_detection = 15  # 法律事件检测模型
    # Model type for text detection and recognition.
    text_detect = 16  # 文字检测+识别
    # Model type for sentence embedding.
    sentence_embedding = 17  # 语义表示
    # Model type for general object detection using Vision Transformer (ViT).
    vit_detection = 18  # ViT_通用目标检测
    # Model type for general semantic segmentation using Vision Transformer (ViT).
    vit_segmentation = 19  # ViT_通用语义分割
    # Model type for face detection.
    face_detection = 20  # 人脸检测
    # Model type for car detection.
    car_detection = 21  # 车辆检测
    # Model type for image classification.
    image_classify = 22  # 图片分类
    # Model type for text generation.
    text_generation = 23  # 文本生成
    # Model type for text classification.
    text_classify = 24  # 文本分类
    # Model type for text similarity.
    text_similarity = 25  # 文本相似度
    # Model type for object detection.
    object_detection = 26  # 目标检测
    # Model type for semantic segmentation.
    semantic_segmentation = 27  # 语义分割
    # Model type for Optical Character Recognition (OCR).
    ocr = 28  # 文字识别
    # Model type for text detection.
    text_detection = 29  # 文字检测
    # Model type for speech classification.
    speech_classify = 30  # 语音分类
    # Model type for speech recognition.
    speech_recognition = 31  # 语音识别
    # Model type for entity recognition.
    entity_recognition = 32  # 实体识别
    # Base model type for image classification.
    image_classify_base = 33  # 图片分类base
    # Large model type for image classification.
    image_classify_large = 34  # 图片分类large
    # Model type for road surface defect detection.
    defect_detection = 35  # 路面瑕疵检测
    # Model type for text-to-image generation (GLIDE version).
    txt2image_glide = 36  # 以文生图，GLIDE版
    # Model type for text-to-image generation (DIFFUSION version).
    txt2image_diffusion = 37  # 以文生图，DIFFUSION版
    # Model type for text-to-image generation (WUKONG version).
    txt2image_wukong = 38  # 以文生图，WUKONF版
    # Model type for multimodal emotion recognition (BERT-based for text).
    multimodal_emotion_recognition_bert = 39  # 多模态情感识别-文字
    # Model type for multimodal emotion recognition (Transformer-based for speech).
    multimodal_emotion_recognition_transformer = 40  # 多模态情感识别-语音
    # Model type for text-to-image generation (Wukong Dreambooth version).
    txt2image_wukong_dreambooth = 41  # 以文生图 wukong dreambooth
    # Model type for OCR on ID cards.
    ocr_idcard = 42  # 文字检测-身份证
    # Model type for OCR on financial documents (invoices).
    ocr_finance = 43  # 文字检测-发票
    # Model type for combined object detection and identification.
    detect_identify = 44  # 目标检测+识别
    # Model type for NLP relation extraction.
    nlp_relation_extraction = 45  # 实体关系抽取
    # Model type for ViT-based object detection with bounding boxes.
    vit_detection_bounding_box = 46  # 目标检测矩形框
    # Model type for rich document information extraction.
    rich_document_infomation_extraction = 47  # 富文档信息抽取
    # Model type for visual question answering.
    visual_question_answer = 48  # 图文问答
    # Model type for multi-round visual question answering.
    multiple_rounds_of_visual_question_answer = 49  # 图文问答_多轮
    # Model type for text question answering.
    text_question_answer = 50  # 文本问答
    # Model type for multi-round text question answering.
    multiple_rounds_of_text_question_answer = 51  # 文本问答_多轮
    # Model type for summary extraction.
    summary_extraction = 52  # 摘要提取
    # Model type for translation.
    translation = 53  # 翻译
    # Model type for text question answering (Qwen version).
    text_question_answer_qwen = 54  # 文本问答_千问
    # Model type for Qwen 2.0 text question answering (0.5B parameters).
    text_question_answer_qwen_2_0_5B = 55  # 千问2.0 文本问答0.5B
    # Model type for Qwen 2.0 text question answering (1.8B parameters).
    text_question_answer_qwen_2_1_8B = 56  # 千问2.0 文本问答1.8B
    # Model type for Gemma text question answering (2B parameters).
    text_question_answer_gemma_2B = 57  # gemma 文本问答2B
    # Model type for Gemma text question answering (7B parameters).
    text_question_answer_gemma_7B = 58  # gemma 文本问答7B
    # Model type for Qwen 2.0 text question answering (4B parameters).
    text_question_answer_qwen_2_4B = 59  # 千问2.0 文本问答4B
    # Model type for Qwen 2.0 text question answering (7B parameters).
    text_question_answer_qwen_2_7B = 60  # 千问2.0 文本问答7B
    # Model type for Llama2 text question answering (7B parameters).
    text_question_answer_llama2_7B = 61  # llama2 文本问答7B
    # Model type for ChatGLM3 text question answering (6B parameters).
    text_question_answer_chatglm3_6B = 62  # chatglm3 文本问答6B
    # Model type for Llama3 text question answering (8B parameters).
    text_question_answer_llama3_8B = 63  # llama3 文本问答8B
    # Model type for UNet Bowl segmentation.
    unet_bowl = 64
    # Model type for YOLOv5 object detection.
    yolov5_detect = 65
    # Model type for DeepLabV3 Cityscapes segmentation.
    deeplabv3_cityscapes = 66
    # Model type for Taichu 2.0 text question answering (2B parameters).
    text_question_answer_taichu_2_0_2B = 67
    # Model type for Taichu 2.0 text question answering (14B parameters).
    text_question_answer_taichu_2_0_14B = 68
    # Model type for Taichu 2.0 text question answering (20B parameters).
    text_question_answer_taichu_2_0_20B = 69
    # Model type for Taichu 2.0 text question answering (32B parameters).
    text_question_answer_taichu_2_0_32B = 70
    # Model type for Taichu 2.0 visual question answering (2B parameters).
    visual_question_answer_taichu_2_0_2B = 71  # taichu2.0 图文问答2B
    # Model type for Taichu 2.0 visual question answering (8B parameters).
    visual_question_answer_taichu_2_0_8B = 72  # taichu2.0 图文问答8B
    # Model type for Qwen 2.5 text question answering (3B parameters).
    text_question_answer_qwen_2_5_3B = 73  # 千问2.5 文本问答3B
    # Model type for Qwen 2.5 text question answering (14B parameters).
    text_question_answer_qwen_2_5_14B = 74  # 千问2.5 文本问答14B
    # Model type for Qwen 2.5 text question answering (32B parameters).
    text_question_answer_qwen_2_5_32B = 75  # 千问2.5 文本问答32B
    # Model type for Qwen 2.5 text question answering (72B parameters).
    text_question_answer_qwen_2_5_72B = 76  # 千问2.5 文本问答72B


# 服务状态枚举 running/deploying/concerning/failed/stopped/finished
# Defines an enumeration for various service statuses.
@unique
class ServiceStatus(TEnum):
    # Indicates that the service is currently running.
    running = 'running'
    # Indicates that the service is currently being deployed.
    deploying = 'deploying'
    # Indicates that the service deployment or operation has failed.
    failed = 'failed'  # 失败
    # Indicates that the service is in a concerning state, possibly with warnings or alerts.
    concerning = 'concerning'  # 告警中
    # Indicates that the service is stopped.
    stopped = 'stopped'
    # Indicates that the service is in the process of stopping.
    stopping = 'stopping'
    # Indicates that the service has finished its operation.
    finished = 'finished'
    # Indicates that the service is in the process of being deleted.
    deleting = 'deleting'
    # Indicates that the service is currently being upgraded.
    upgrading = 'upgrading'  # 升级中

    # 服务状态是否可以预测
    # Checks if the service status allows for prediction.
    # Args:
    #   status: The current status of the service (string or ServiceStatus enum).
    # Returns:
    #   True if the service can predict, False otherwise.
    @staticmethod
    def can_predict(status):
        if isinstance(status, str):
            return status in [ServiceStatus.running.value, ServiceStatus.upgrading.value]

        if isinstance(status, ServiceStatus):
            return status in [ServiceStatus.running, ServiceStatus.upgrading]

        return False

    @staticmethod
    # Checks if the service can be started from its current status.
    # Args:
    #   status: The current status of the service (string or ServiceStatus enum).
    # Returns:
    #   True if the service can be started, False otherwise.
    def can_start(status):
        if isinstance(status, str):
            return status in [ServiceStatus.stopped.value]

        if isinstance(status, ServiceStatus):
            return status in [ServiceStatus.stopped]

        return False

    @staticmethod
    # Checks if the service can be stopped from its current status.
    # Args:
    #   status: The current status of the service (string or ServiceStatus enum).
    # Returns:
    #   True if the service can be stopped, False otherwise.
    def can_stop(status):
        allowed_status = [
            ServiceStatus.running.value,
            ServiceStatus.upgrading.value,
            ServiceStatus.failed.value,
            ServiceStatus.deploying.value,
            ServiceStatus.concerning.value,
        ]

        if isinstance(status, str):
            return status in allowed_status

        if isinstance(status, ServiceStatus):
            return status in allowed_status

        return False

    @staticmethod
    # Checks if the service can be updated from its current status.
    # Args:
    #   status: The current status of the service (string or ServiceStatus enum).
    # Returns:
    #   True if the service can be updated, False otherwise.
    def can_update(status):
        if isinstance(status, str):
            return status in [
                ServiceStatus.running.value,
                ServiceStatus.stopped.value,
                ServiceStatus.upgrading.value,
            ]

        if isinstance(status, ServiceStatus):
            return status in [
                ServiceStatus.running,
                ServiceStatus.stopped,
                ServiceStatus.upgrading,
            ]

        return False

    @staticmethod
    # Checks if the service can be upgraded from its current status.
    # Args:
    #   status: The current status of the service (string or ServiceStatus enum).
    # Returns:
    #   True if the service can be upgraded, False otherwise.
    def can_upgrade(status):
        if isinstance(status, str):
            return status in [ServiceStatus.running.value, ServiceStatus.upgrading.value]

        if isinstance(status, ServiceStatus):
            return status in [ServiceStatus.running, ServiceStatus.upgrading]

        return False


# 模型场景
# Defines an enumeration for different model scenes.
class EnumModelScene(TEnum):
    # Represents an unknown or other model scene.
    other = defaultInt
    # Represents the AI model scene.
    ai = 1


# 模型维度
# Defines an enumeration for different model dimensions.
class EnumModelDimension(TEnum):
    # Represents an unknown or other model dimension.
    other = defaultInt
    # Represents the omni-modal dimension.
    omni = 1
    # Represents the vision dimension.
    vision = 2
    # Represents the natural language processing (NLP) dimension.
    nlp = 3
    # Represents the voice dimension.
    voice = 4
    # Represents the Optical Character Recognition (OCR) dimension.
    ocr = 5
    # Represents a chosen or selected dimension.
    chosen = 6
    # Represents the digital human dimension.
    digital_human = 7  # 虚拟人


# Configuration for model dimensions, frameworks, and scenes.
MODEL_DIMENSION_CFGS = {
    'dimension': [
        {
            'class': EnumModelDimension.nlp.name,
            'cn_name': 'NLP(自然语言处理)',
            'subclass': [
                {'class': EnumModelType.text_generation.name, 'cn_name': '文本生成'},
                {
                    'class': EnumModelType.entity_extraction.name,
                    'cn_name': '实体关系抽取',
                },
                {
                    'class': EnumModelType.text_classify.name,
                    'cn_name': '文本分类',
                },
                {
                    'class': EnumModelType.text_similarity.name,
                    'cn_name': '文本相似度',
                },
                {
                    'class': EnumModelType.entity_recognition.name,
                    'cn_name': '实体识别',
                },
            ],
        },
        {
            'class': EnumModelDimension.omni.name,
            'cn_name': '多模态',
            'subclass': [
                {'class': EnumModelType.image_caption.name, 'cn_name': '以图生文'},
                {'class': EnumModelType.txt2image.name, 'cn_name': '以文生图'},
                {'class': EnumModelType.vqa.name, 'cn_name': '视觉问答'},
                {
                    'class': EnumModelType.retrieval.name,
                    'cn_name': '图文检索',
                },
                # {
                #     'class': EnumModelType.multiple_rounds_of_visual_question_answer.name,
                #     'cn_name': '图文问答',
                # },
                {
                    'class': EnumModelType.visual_question_answer.name,
                    'cn_name': '图文问答',
                },
            ],
        },
        {
            'class': EnumModelDimension.voice.name,
            'cn_name': '语音',
            'subclass': [
                {
                    'class': EnumModelType.speech_recognition.name,
                    'cn_name': '语音识别',
                },
                {
                    'class': EnumModelType.speech_classify.name,
                    'cn_name': '语音分类',
                },
            ],
        },
        {
            'class': EnumModelDimension.vision.name,
            'cn_name': '视觉',
            'subclass': [
                {
                    'class': EnumModelType.image_classify.name,
                    'cn_name': '图片分类',
                },
                {
                    'class': EnumModelType.object_detection.name,
                    'cn_name': '目标检测',
                },
                {
                    'class': EnumModelType.semantic_segmentation.name,
                    'cn_name': '语义分割',
                },
                {
                    'class': EnumModelType.ocr.name,
                    'cn_name': '文字识别',
                },
                {
                    'class': EnumModelType.text_detection.name,
                    'cn_name': '文字检测',
                },
            ],
        },
        {'class': EnumModelDimension.other.name, 'cn_name': '其他'},
    ],
    'framework': [
        {'class': EnumFramework.pytorch.name, 'cn_name': 'PyTorch'},
        {'class': EnumFramework.mindspore.name, 'cn_name': 'MindSpore'},
        {'class': EnumFramework.other.name, 'cn_name': '其他'},
    ],
    'scene': [
        {
            'class': EnumModelDimension.omni.name,
            'cn_name': '多模态',
            'subclass': [{'class': EnumModelScene.ai.name, 'cn_name': 'AI创作'}],
        }
    ],
}


# Defines an enumeration for various model export statuses.
class EnumModelExportStatus(TEnum):
    # Indicates that the model is currently being exported.
    exporting = 'exporting'  # 导出中
    # Indicates that the model export was successful.
    success = 'success'  # 导出成功
    # Indicates that the model export failed.
    failed = 'failed'  # 导出失败
    # Indicates that the model export has expired.
    expired = 'expired'  # 已过期
    # Indicates that the model export is waiting in a queue.
    waiting = 'waiting'  # 排队中


# Defines an enumeration for various notebook statuses.
@unique
class EnumNotebookStatus(TEnum):
    # Indicates that the notebook is currently running.
    running = 'running'
    # Indicates that the notebook is currently being deployed.
    deploying = 'deploying'
    # Indicates that the notebook operation has failed.
    failed = 'failed'  # 失败
    # Indicates that the notebook is stopped.
    stopped = 'stopped'
    # Indicates that the notebook is in the process of stopping.
    stopping = 'stopping'
    # Indicates that the notebook is in the process of being deleted.
    deleting = 'deleting'
    # Indicates that the notebook is currently being saved or built.
    saving = 'saving'  # 构建中

    @staticmethod
    # Checks if the notebook can be started from its current status.
    # Args:
    #   status: The current status of the notebook (string or EnumNotebookStatus enum).
    # Returns:
    #   True if the notebook can be started, False otherwise.
    def can_start(status):
        if isinstance(status, str):
            return status in [EnumNotebookStatus.stopped.value]

        if isinstance(status, EnumNotebookStatus):
            return status in [EnumNotebookStatus.stopped]

        return False

    @staticmethod
    # Checks if the notebook can be stopped from its current status.
    # Args:
    #   status: The current status of the notebook (string or EnumNotebookStatus enum).
    # Returns:
    #   True if the notebook can be stopped, False otherwise.
    def can_stop(status):
        if isinstance(status, str):
            return status in [
                EnumNotebookStatus.running.value,
                EnumNotebookStatus.deploying.value,
                EnumNotebookStatus.failed.value,
                EnumNotebookStatus.stopped.value,
            ]

        if isinstance(status, EnumNotebookStatus):
            return status in [
                ServiceStatus.running,
                EnumNotebookStatus.deploying,
                EnumNotebookStatus.failed,
                EnumNotebookStatus.stopped,
            ]

        return False


# Defines an enumeration for various Docker commit statuses.
@unique
class EnumDockerCommitStatus(TEnum):
    # Indicates that the Docker commit is pending.
    pending = 'pending'
    # Indicates that the Docker commit is currently running.
    running = 'running'
    # Indicates that the Docker commit failed.
    failed = 'failed'
    # Indicates that the Docker commit succeeded.
    succeeded = 'succeeded'